package com.patsnap.data.npd.dw.etl.job.ods.job;

import com.google.common.collect.ImmutableMap;
import com.patsnap.data.npd.dw.etl.job.ods.serialization.PreOdsInfo;
import com.patsnap.data.npd.dw.etl.job.ods.serialization.PreOdsInfoDeserializationSchema;
import com.patsnap.one.etl.flink.job.AbstractFlinkJobLauncher;
import com.patsnap.one.etl.tool.datasource.DataSourceAssembler;
import com.patsnap.one.etl.tool.kafka.KafkaConfig;
import com.patsnap.one.etl.tool.kafka.KafkaHelper;
import com.patsnap.one.etl.version.Version;
import org.apache.flink.api.common.eventtime.WatermarkStrategy;
import org.apache.flink.api.java.utils.ParameterTool;
import org.apache.flink.connector.kafka.source.KafkaSource;
import org.apache.flink.streaming.api.CheckpointingMode;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.CheckpointConfig;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.kafka.clients.consumer.ConsumerConfig;

import java.util.Map;

import static com.patsnap.one.etl.constant.Constants.LAUNCH_ARGS;
import static com.patsnap.one.etl.constant.Constants.LAUNCH_CLASS;
import static com.patsnap.one.etl.version.Version.EASY_ETL_COMPONENT_VERSION_KEY;

public abstract class OdsFlinkJobLauncher extends AbstractFlinkJobLauncher {

    protected static final String CHECK_POINT_IS = "is_checkpoint";

    protected SingleOutputStreamOperator<PreOdsInfo> generateSource(StreamExecutionEnvironment env, ParameterTool parameterTool, String kafkaSourceName) {
        Map<String, KafkaConfig> kafkaConfigMap = DataSourceAssembler.buildKafkaConfigsFromParameterTool(parameterTool);
        KafkaConfig kafkaConfig = kafkaConfigMap.get(kafkaSourceName);
        KafkaSource<PreOdsInfo> kafkaSource = KafkaHelper.<PreOdsInfo>builder(kafkaConfig)
                .setDeserializer(new PreOdsInfoDeserializationSchema())
                .setProperty(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, parameterTool.get("is_kafka_auto_commit", "false"))
                .build();

        return env.fromSource(kafkaSource, WatermarkStrategy.noWatermarks(), kafkaConfig.getTopics());
    }

    protected StreamExecutionEnvironment getExecutionEnvironment(int sourceParallelism, ParameterTool parameterTool, String[] args, String launchClass) {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        ParameterTool launcherParams = ParameterTool.fromMap(ImmutableMap.of(LAUNCH_ARGS, String.join(" ", args), LAUNCH_CLASS, launchClass, EASY_ETL_COMPONENT_VERSION_KEY, Version.EASY_ETL_COMPONENT_VERSION_VALUE));
        env.getConfig().setGlobalJobParameters(parameterTool.mergeWith(launcherParams));
        env.setParallelism(sourceParallelism);
        return env;
    }

    protected void enableCheckpointing(StreamExecutionEnvironment env, ParameterTool parameterTool) {
        String checkpointInterval = parameterTool.get("checkpoint_interval_sec", "10");
        String checkpointTimeout = parameterTool.get("checkpoint_timeout_sec", "600");
        String miniPauseBetween = parameterTool.get("mini_pause_between_sec", "4");
        String enableUnalignedCheckpoints = parameterTool.get("unalign_checkpoint", "false");
        env.enableCheckpointing(Long.parseLong(checkpointInterval) * 1000L);
        CheckpointConfig checkpointConfig = env.getCheckpointConfig();
        checkpointConfig.setCheckpointingMode(CheckpointingMode.AT_LEAST_ONCE);
        checkpointConfig.setExternalizedCheckpointCleanup(CheckpointConfig.ExternalizedCheckpointCleanup.RETAIN_ON_CANCELLATION);
        if (Boolean.TRUE.equals(Boolean.valueOf(enableUnalignedCheckpoints))) {
            checkpointConfig.enableUnalignedCheckpoints();
        }

        checkpointConfig.setCheckpointTimeout(Long.parseLong(checkpointTimeout) * 1000L);
        checkpointConfig.setMinPauseBetweenCheckpoints(Long.parseLong(miniPauseBetween) * 1000L);
    }
}
